#!/usr/bin/env python3
import rospy
import copy
from control_msgs.msg import FollowJointTrajectoryActionGoal
from trajectory_msgs.msg import JointTrajectoryPoint

class HostToFollowRelay:
    def __init__(self):
        # params (可在 launch/rosparam 中覆写)
        self.host_topic = rospy.get_param("~host_topic",
                                          "/robot_jiazhao_control/follow_joint_trajectory/goal")
        self.follow_topics = rospy.get_param("~follow_topics", [
            "/follow1_joint/follow_joint_trajectory/goal",
            "/follow2_joint/follow_joint_trajectory/goal"
        ])
        self.host_joint = rospy.get_param("~host_joint", "host_joint")
        self.follow_joints = rospy.get_param("~follow_joints", ["follow1_joint", "follow2_joint"])
        self.multipliers = rospy.get_param("~multipliers", [1.0, -1.0])  # 默认第二指反向
        self.offsets = rospy.get_param("~offsets", [0.0, 0.0])

        assert len(self.follow_topics) == len(self.follow_joints) == len(self.multipliers) == len(self.offsets), \
            "follow_topics, follow_joints, multipliers, offsets must have same length"

        # publishers for each follow topic (same message type as host goal)
        self.pubs = [rospy.Publisher(t, FollowJointTrajectoryActionGoal, queue_size=10) for t in self.follow_topics]

        # subscribe host goal
        rospy.Subscriber(self.host_topic, FollowJointTrajectoryActionGoal, self.cb, queue_size=10)

        rospy.loginfo("HostToFollowRelay started. host_topic=%s -> %s", self.host_topic, ", ".join(self.follow_topics))

    def cb(self, msg: FollowJointTrajectoryActionGoal):
        # 确认 host_joint 在 incoming joint_names 中
        if not msg.goal or not msg.goal.trajectory or not msg.goal.trajectory.joint_names:
            rospy.logwarn("received goal has no trajectory/joint_names")
            return

        try:
            idx = msg.goal.trajectory.joint_names.index(self.host_joint)
        except ValueError:
            rospy.logwarn_throttle(5.0, "host_joint '%s' not found in incoming joint_names: %s",
                                   self.host_joint, msg.goal.trajectory.joint_names)
            return

        # 对每个 follow controller，构造新的 FollowJointTrajectoryActionGoal 并发布
        for pub, fj, mul, off in zip(self.pubs, self.follow_joints, self.multipliers, self.offsets):
            new_msg = copy.deepcopy(msg)  # 复制 header, goal_id, 等
            # 只保留目标 follow joint 的名字
            new_msg.goal.trajectory.joint_names = [fj]
            new_points = []
            for pt in msg.goal.trajectory.points:
                # positions/velocities/accelerations 可能为空，先处理
                pos = []
                vel = []
                acc = []
                if pt.positions:
                    val = pt.positions[idx] * mul + off
                    pos = [val]
                if pt.velocities:
                    vel = [pt.velocities[idx] * mul]
                if pt.accelerations:
                    acc = [pt.accelerations[idx] * mul]

                new_pt = JointTrajectoryPoint()
                new_pt.positions = pos
                if vel:
                    new_pt.velocities = vel
                if acc:
                    new_pt.accelerations = acc
                # effort 留空（通常没必要），保留时间戳
                new_pt.time_from_start = pt.time_from_start
                new_points.append(new_pt)

            new_msg.goal.trajectory.points = new_points

            # publish
            pub.publish(new_msg)
            rospy.logdebug("Relayed host -> %s (joint %s)", pub.name, fj)


if __name__ == "__main__":
    rospy.init_node("host_to_follow_relay")
    node = HostToFollowRelay()
    rospy.spin()
